import os
from utils.common.logger import log
from snownlp import SnowNLP, sentiment

pre_trained_model = os.path.dirname(__file__) + "/data/sentiment.marshal"


class SentimentAnalyzer:
    '''
    对文本进行情感分析
    '''

    def __init__(self, model_path: str = pre_trained_model):
        sentiment.load(model_path)
        log.logger.info("SentimentAnalyzer initialized...")

    def train(self, neg_path: str, pos_path: str):
        '''
        用户提供neg.txt和pos.txt，训练模型并输出
        '''
        sentiment.train(neg_path, pos_path)
        sentiment.save("sentiment.marshal")

    def senti_news(self, news_title: str, news_content: str):
        '''
        调用序列化的模型 给新闻的标题和内容分别打标签

        正面：1；负面：-1；中立：0
        '''

        title_senti = SnowNLP(news_title).sentiments
        content_senti = SnowNLP(news_content).sentiments

        if title_senti < 0.3 or content_senti < 0.1 or (0.7*title_senti + 0.3*content_senti) < 0.3:
            return -1
        elif title_senti > 0.7 or content_senti > 0.9 or (0.7*title_senti + 0.3*content_senti) > 0.7:
            return 1

        return 0
